import torch
import torchsparse
import torchsparse.backend
import torch.utils.data

from torchsparse import SparseTensor
from torchsparse import nn as spnn

# Create a dummy sparse matrix
num_points = 10

coords = torch.randint(4, size=(num_points, 4)).int()
coords[:, 0] = 0
features = torch.rand(num_points)[:, None]

#dummy = ME.SparseTensor(coordinates=coords, features=features)
dummy = SparseTensor(features, coords, stride=2).to('cuda')
print(f"{dummy.stride=}")
# Create a strided convolution, apply to the dummy input
# conv = ME.MinkowskiConvolution(in_channels=1, out_channels=4, kernel_size=[2, 2, 2], stride=[2, 2, 2], dilation=[1, 1, 1], dimension=3)
conv = spnn.Conv3d(in_channels=1, out_channels=4, kernel_size=[2, 2, 2], stride=[2, 2, 2], dilation=[1,1,1],transposed=False).to('cuda')
x = conv(dummy)
print(f'{x.feats=} {x.coords=} {x.stride=}')


# Create a strided transpose convolution, apply
#conv_t = ME.MinkowskiConvolutionTranspose(in_channels=4, out_channels=4, kernel_size=[2, 2, 2], stride=[2, 2, 2], dilation=[1, 1, 1], dimension=3)
#x = conv_t(x)

conv_t = spnn.Conv3d(in_channels=4, out_channels=4, kernel_size=2, stride=2, dilation=[1,1,1], transposed=True).to('cuda')
x = conv_t(x)
print(f'{x.feats=} {x.coords=} {x.stride=}')